#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <errno.h>
#include <unistd.h>
#include <time.h>

// Linux defaults to an easier to use format
// http://en.wikipedia.org/wiki/Tcphdr
// we want to use the legacy one that works on OS X too.
//#define __USE_BSD
#define __FAVOR_BSD

#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <sys/types.h>
#include <sys/socket.h>

// ts.tv_sec
// ts.tv_nsec
long long get_time_ns() {
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &ts);
  long long time = ((long long)ts.tv_sec) * 1000000000 + ts.tv_nsec;
  return time;
}

// Checksum code from http://www.enderunix.org/docs/en/rawipspoof/
u_short ip_checksum(u_short* addr, int len) {
  int nleft = len;
  int sum = 0;
  u_short* w = addr;
  u_short answer = 0;
  while (nleft > 1) {
    sum += *w++;
    nleft -= 2;  
  }
  if (nleft == 1) {
    *(u_char*)(&answer) = *(u_char*)w;
    sum += answer;
  }
  sum = (sum >> 16) + (sum & 0xFFFF);
  sum += (sum >> 16);
  answer = ~sum;
  return answer;
}
struct pseudo_udp {
  struct in_addr src;
  struct in_addr dst;
  u_char pad;
  u_char proto;
  u_short udp_len;
  struct udphdr udp;
};
u_short udp_checksum(in_addr_t src, in_addr_t dst, u_short* addr, int len) {
  struct pseudo_udp buf;
  memset(&buf, 0, sizeof(buf));
  buf.src.s_addr = src;
  buf.dst.s_addr = dst;
  buf.pad = 0;
  buf.proto = IPPROTO_UDP;
  buf.udp_len = htons(len);
  memcpy(&(buf.udp), addr, len);
  return ip_checksum((u_short*)&buf, 12 + len);
}
struct pseudo_tcp {
  struct in_addr src;
  struct in_addr dst;
  u_char pad;
  u_char proto;
  u_short tcp_len;
  struct tcphdr tcp;
};
u_short tcp_checksum(in_addr_t src, in_addr_t dst, u_short* addr, int len) {
  struct pseudo_tcp buf;
  memset(&buf, 0, sizeof(buf));
  buf.src.s_addr = src;
  buf.dst.s_addr = dst;
  buf.pad = 0;
  buf.proto = IPPROTO_TCP;
  buf.tcp_len = htons(len);
  memcpy(&(buf.tcp), addr, len);
  return ip_checksum((u_short*)&buf, 12 + len);
}

# define DEFAULT_IP_HEADER_LENGTH 20 // bytes
u_char* spoof_ip(in_addr_t src, in_addr_t dst, u_int16_t ip_id, u_int16_t protocol, u_int16_t len, u_int16_t* out_len) {
  u_char* packet;
  struct ip ip_header;
  u_int16_t total_len = DEFAULT_IP_HEADER_LENGTH + len;
  packet = (u_char*)malloc(total_len);

  ip_header.ip_hl = 0x5;                                                      // [4]  header length = 20 bytes
  ip_header.ip_v = 0x4;                                                       // [4]  ip version = 4
  ip_header.ip_tos = 0x0;                                                     // [8]  type of service = normal
  ip_header.ip_len = total_len;                                               // [16] length
  ip_header.ip_id = htons(ip_id);                                             // [16] unique packet id
  
// [16] offset, don't fragment
#ifdef __APPLE__
  ip_header.ip_off = IP_DF;
#else
  ip_header.ip_off = htons(IP_DF); 
#endif
  
  ip_header.ip_ttl = 64;                                                      // [8]  TTL
  ip_header.ip_p = protocol;                                                  // [8]  layer 3 protocol
  ip_header.ip_sum = 0x0;                                                     // [16] IP header checksum, will set this later
  ip_header.ip_src.s_addr = src;                                              // [16] source address
  ip_header.ip_dst.s_addr = dst;                                              // [16] destination address
  ip_header.ip_sum = ip_checksum((u_short*)&ip_header, sizeof(ip_header));
  memcpy(packet, &ip_header, sizeof(ip_header));
  *out_len = total_len;
  return packet;
}

u_char* spoof_tcp(in_addr_t src, u_int16_t src_port,
                  in_addr_t dst, u_int16_t dst_port,
                  u_int16_t ip_id,
                  u_int32_t sequence_number,
                  u_int32_t acknowledgement_number,
                  u_int8_t flags,                                             // TCP packet type
                  u_int16_t window_size,
                  u_char* data, u_int16_t len, u_int16_t* out_len) {
  u_char* packet;
  struct tcphdr tcp_header;
  
  u_int16_t tcp_len = sizeof(tcp_header) + len;

  tcp_header.th_sport = htons(src_port);                                      // Source port
  tcp_header.th_dport = htons(dst_port);                                      // Destination port
  tcp_header.th_seq = htonl(sequence_number);
  tcp_header.th_ack = htonl(acknowledgement_number);
  tcp_header.th_off = sizeof(struct tcphdr) / 4;                              // For lack of a better value
  tcp_header.th_flags = flags;
  tcp_header.th_win = htons(window_size);
  tcp_header.th_sum = 0x0;                                                    // Checksum should be 0 for checksum calculation
  tcp_header.th_sum = tcp_checksum(src, dst, (u_short*)&tcp_header, sizeof(tcp_header));
  
  packet = spoof_ip(src, dst, ip_id, IPPROTO_TCP, tcp_len, out_len);          // IP
  memcpy(packet + DEFAULT_IP_HEADER_LENGTH, &tcp_header, sizeof(tcp_header)); // TCP header
  memcpy(packet + DEFAULT_IP_HEADER_LENGTH + sizeof(tcp_header), data, len);  // Payload
  return packet;
}

u_char* spoof_udp(in_addr_t src, u_int16_t src_port,
                  in_addr_t dst, u_int16_t dst_port,
                  u_int16_t ip_id,
                  u_char* data,
                  u_int16_t len,
                  u_int16_t* out_len) {
  u_char* packet;
  struct udphdr udp_header;
  u_int16_t udp_len = sizeof(udp_header) + len;

  udp_header.uh_sport = htons(src_port);                                      // Source port
  udp_header.uh_dport = htons(dst_port);                                      // Destination port
  udp_header.uh_ulen = htons(sizeof(udp_header) + len);                       // Length of UDP packet (header + data)
  udp_header.uh_sum = 0x0;                                                    // Checksum should be 0 for checksum calculation
  udp_header.uh_sum = udp_checksum(src, dst, (u_short*)&udp_header, sizeof(udp_header));
  
  packet = spoof_ip(src, dst, ip_id, IPPROTO_UDP, udp_len, out_len);          // IP
  memcpy(packet + DEFAULT_IP_HEADER_LENGTH, &udp_header, sizeof(udp_header)); // UDP header
  memcpy(packet + DEFAULT_IP_HEADER_LENGTH + sizeof(udp_header), data, len);  // Payload
  return packet;
}

void spoof_loop(int socket_descriptor,
                u_int8_t protocol,
                int repeat,
		float offset,
                in_addr_t src, u_int16_t src_port,
                in_addr_t dst, u_int16_t dst_port,
                u_int16_t size) {
  u_char* packet = NULL;
  int i;
  
  u_int8_t packet_data[size];
  u_int16_t len;
  u_int16_t seq = 6231;

  struct sockaddr_in sin;
  memset(&sin, 0, sizeof(sin));
  sin.sin_family = AF_INET;
  sin.sin_addr.s_addr = src;

  // Put current time and add to packet.
  long ns_offset = offset * 1000000000;
  u_int64_t time = get_time_ns() + ns_offset;
  memcpy(packet_data, &time, sizeof(time));
  gethostname((void*)packet_data+sizeof(time), size-sizeof(time));
  
  for (i = 0; i < repeat; i++) {
    if (protocol == IPPROTO_UDP) {
      packet = spoof_udp(
          src, src_port, dst, dst_port,
          1234 + i,
          (u_char*)&packet_data,
          sizeof(packet_data),
          &len
        );
    }
    else if (protocol == IPPROTO_TCP) {
      packet = spoof_tcp(
          src, src_port, dst, dst_port,
          1234 + i, // ip_id
          seq,      // seq
          7814 + i, // ack
          TH_ACK,
          272,      // window size
          (u_char*)&packet_data,
          sizeof(packet_data),
          &len
        );
      seq += size;
    }
    
    if (sendto(socket_descriptor, packet, len, 0, (struct sockaddr*)&sin, sizeof(struct sockaddr)) < 0) {
      fprintf(stderr, "Failed to send packet %d: %s (%d)\n", i, strerror(errno), errno);
    }
    if (packet != NULL) {
      free(packet);
      packet = NULL;
      
      // Sleep for a little while. CPU is much faster than network.
      usleep(1000);
    }
  }
}

int open_raw_socket() {
  const int on = 1;
  int socket_descriptor;
  if ((socket_descriptor = socket(AF_INET, SOCK_RAW, IPPROTO_RAW)) < 0) {
    fprintf(stderr, "Failed to create raw socket. Are you root?\n");
    exit(EXIT_FAILURE);
  }
  if (setsockopt(socket_descriptor, IPPROTO_IP, IP_HDRINCL, &on, sizeof(on)) < 0) {
    fprintf(stderr, "Failed to set socket options.\n");
    exit(EXIT_FAILURE);
  }
  return socket_descriptor;
}

int main(int argc, char** argv) {
  in_addr_t src;
  in_addr_t dst;
  int src_port;
  int dst_port;
  u_int8_t protocol = IPPROTO_UDP;
  int count = 0;
  int payload_size = 1448;
  int socket_descriptor;
  float offset;
  
  if (argc < 7) {
    fprintf(stderr, "Usage: ./tx protocol count time_offset src_ip src_port dst_ip dst_port [payload_size=1448]\n");
    exit(EXIT_FAILURE);
  }
  
  socket_descriptor = open_raw_socket();
  
  if (strcasecmp(argv[1], "TCP") == 0) {
    protocol = IPPROTO_TCP;
  }
  else if (strcasecmp(argv[1], "UDP") == 0) {
    protocol = IPPROTO_UDP;
  }
  sscanf(argv[2], "%d", &count);
  sscanf(argv[3], "%f", &offset);
  
  src = inet_addr(argv[4]); // 50.38.0.159
  sscanf(argv[5], "%d", &src_port);
  
  dst = inet_addr(argv[6]); // 192.168.1.8
  sscanf(argv[7], "%d", &dst_port);
  
  if (argc > 8) {
    sscanf(argv[8], "%d", &payload_size);
  }
  
  spoof_loop(socket_descriptor,
             protocol,
             count,
	     offset,
             src,
             (u_int16_t)src_port,
             dst,
             (u_int16_t)dst_port,
             (u_int16_t)payload_size);
  
  return 0;
}

